package com.efeichong.common.config.flowLimit;

import com.efeichong.util.FastJpaSpringUtils;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

import java.time.Instant;
import java.util.List;

/**
 * @author lxk
 * @date 2022/8/22
 * @description 限流处理(令牌)
 */
@Slf4j
public class TokenFlowLimitHandler {

    private static final String LUA_SCRIPT = "local tokens_key = KEYS[1] -- request_rate_limiter.${id}.tokens 令牌桶剩余令牌数的KEY值\n" +
            "local timestamp_key = KEYS[2] -- 令牌桶最后填充令牌时间的KEY值\n" +
            "--redis.log(redis.LOG_WARNING, \"tokens_key \" .. tokens_key)\n" +
            " \n" +
            "local rate = tonumber(ARGV[1]) --  令令牌桶填充平均速率\n" +
            "local capacity = tonumber(ARGV[2]) --  令牌桶上限\n" +
            "local now = tonumber(ARGV[3]) -- 得到从 1970-01-01 00:00:00 开始的秒数\n" +
            "local requested = tonumber(ARGV[4]) -- 消耗令牌数量，默认 1\n" +
            " \n" +
            "local fill_time = capacity/rate -- 计算令牌桶填充满令牌需要多久时间\n" +
            "local ttl = math.floor(fill_time*2) -- *2 保证时间充足\n" +
            " \n" +
            "--redis.log(redis.LOG_WARNING, \"rate \" .. ARGV[1])\n" +
            "--redis.log(redis.LOG_WARNING, \"capacity \" .. ARGV[2])\n" +
            "--redis.log(redis.LOG_WARNING, \"now \" .. ARGV[3])\n" +
            "--redis.log(redis.LOG_WARNING, \"requested \" .. ARGV[4])\n" +
            "--redis.log(redis.LOG_WARNING, \"filltime \" .. fill_time)\n" +
            "--redis.log(redis.LOG_WARNING, \"ttl \" .. ttl)\n" +
            " -- 获得令牌桶剩余令牌数\n" +
            "local last_tokens = tonumber(redis.call(\"get\", tokens_key)) -- 第一次时，没有数值，所以桶时满的\n" +
            "if last_tokens == nil then\n" +
            "  last_tokens = capacity\n" +
            "end\n" +
            "--redis.log(redis.LOG_WARNING, \"last_tokens \" .. last_tokens)\n" +
            " \n" +
            "local last_refreshed = tonumber(redis.call(\"get\", timestamp_key))\n" +
            "-- 令牌桶最后填充令牌时间\n" +
            "if last_refreshed == nil then\n" +
            "  last_refreshed = 0\n" +
            "end\n" +
            "--redis.log(redis.LOG_WARNING, \"last_refreshed \" .. last_refreshed)\n" +
            " \n" +
            "local delta = math.max(0, now-last_refreshed)\n" +
            "-- 获取距离上一次刷新的时间间隔\n" +
            "local filled_tokens = math.min(capacity, last_tokens+(delta*rate))\n" +
            "-- 填充令牌，计算新的令牌桶剩余令牌数 填充不超过令牌桶令牌上限\n" +
            "local allowed = filled_tokens >= requested\n" +
            "local new_tokens = filled_tokens\n" +
            "local allowed_num = 0\n" +
            "if allowed then\n" +
            "-- 若成功，令牌桶剩余令牌数(new_tokens) 减消耗令牌数( requested )，并设置获取成功( allowed_num = 1 )\n" +
            "  new_tokens = filled_tokens - requested\n" +
            "  allowed_num = 1\n" +
            "end\n" +
            " \n" +
            "--redis.log(redis.LOG_WARNING, \"delta \" .. delta)\n" +
            "--redis.log(redis.LOG_WARNING, \"filled_tokens \" .. filled_tokens)\n" +
            "--redis.log(redis.LOG_WARNING, \"allowed_num \" .. allowed_num)\n" +
            "--redis.log(redis.LOG_WARNING, \"new_tokens \" .. new_tokens)\n" +
            " -- 设置令牌桶剩余令牌数( new_tokens ) ，令牌桶最后填充令牌时间(now) ttl是超时时间\n" +
            "redis.call(\"setex\", tokens_key, ttl, new_tokens) -- 失效时间是上面计算的ttl，单位是秒\n" +
            "redis.call(\"setex\", timestamp_key, ttl, now)\n" +
            " -- 返回数组结果，new_tokens表示桶内目前剩余的令牌数\n" +
            " -- 返回的数组，allowed_num=1表示可以访问，0表示不可以访问，桶内没那么多令牌了。\n" +
            "return { allowed_num, new_tokens }  ";
    private static StringRedisTemplate redisTemplate;

    static {
        redisTemplate = FastJpaSpringUtils.getBean(StringRedisTemplate.class);
    }


    public static boolean flow(String key, int capacity) {
        String tokenKey = key + "." + "tokens";
        String timestampKey = key + "." + "timestamp";
        //往桶中放入令牌的速率/秒
        String rate = capacity + "";
        //需要桶中有多少个令牌
        String requested = "1";
        //获取到秒级别，而System.currentTimeMillis()获取到毫秒级
        String now = Instant.now().getEpochSecond() + "";

        RedisScript<List> redisScript = new DefaultRedisScript<>(LUA_SCRIPT, List.class);
        List<Long> execute = redisTemplate.execute(redisScript, Lists.newArrayList(tokenKey, timestampKey), rate, capacity + "", now, requested);
        Long hasAvailable = execute.get(0);
        Long remaining = execute.get(1);
        log.info("hasAvailable={},remaining={}", hasAvailable, remaining);
        return hasAvailable == 1;
    }
}
